pub(crate) mod writer;
pub(crate) mod utils;
pub(crate) mod converter;
pub(crate) mod parser;
pub(crate) mod args;

use std::{env, fs};
use std::path::PathBuf;
use rand::seq::SliceRandom;
use rand::rng;
use log::error;
use writer::create_yolo_dataset;
use crate::utils::image_files;
use crate::args::parse;

// 矩形框坐标结构
#[derive(Debug)]
pub struct Bbox {
    pub x1: i32,
    pub y1: i32,
    pub x2: i32,
    pub y2: i32,
}

// YOLO数据结构
#[derive(Debug)]
pub struct YoloData {
    pub x_center: f32,
    pub y_center: f32,
    pub width: f32,
    pub height: f32,
}

#[derive(Debug)]
enum DataType {
    Test,
    Train,
    Val,
    Unknown
}

impl DataType {
    fn as_str(&self) -> &str {
        match self {
            DataType::Train => "train",
            DataType::Val => "val",
            DataType::Test => "test",
            DataType::Unknown => "unknown",
        }
    }
}


fn main() -> Result<(), Box<dyn std::error::Error>> {

    fs::create_dir_all("logs")?;

    // 初始化 log4rs 配置， 使用log4rs方式
    log4rs::init_file("log4rs.yaml", Default::default())?;

    // 提取参数
    let mut args: Vec<String> = env::args().collect();

    //不传参数，则使用默认参数，测试数据
    if args.len() < 2 {
        args = vec![
            "img2yolo".to_string(),
            "-t".parse().unwrap(),
            "train".to_string(),
            "-s".parse().unwrap(),
            "./assets/images".to_string(), // 图片目录
            "-d".parse().unwrap(),
            "./dataset".to_string(), //目标目录
        ];
    }


    let config = parse(args);

    log::info!("Starting img2yolo...");

    // 获取图片文件
    let mut image_files = image_files(config.src.as_str());
    if image_files.is_empty() {
        error!("No image files found in {}", config.src.as_str());
        std::process::exit(1);
    }

    // 数据集比例, 0-1
    if config.ratio >= 0_f64 && config.ratio < 1_f64 {
        // 随机抽取20%的数据
        let mut rng = rng();
        let sample_size = (image_files.len() as f64 * config.ratio).ceil() as usize;
        image_files.shuffle(&mut rng);
        image_files.truncate(sample_size);
    }
    log::info!("{} files found in {}", image_files.len(), config.src.as_str());

    // 创建数据集目录
    let data_type_folder = config.data_type.as_str();
    let train_path_images = PathBuf::from(config.dst.as_str())
        .join(data_type_folder).join( "images");
    let train_path_labels = PathBuf::from(config.dst.as_str())
        .join(data_type_folder).join( "labels");

    // 创建数据集
    let result = create_yolo_dataset(image_files, config.class_id, train_path_images, train_path_labels);
    if result.is_err() {
        error!("Error: {}", result.err().unwrap());
        std::process::exit(1);
    } else {
        log::info!("{} dataset created successfully", data_type_folder);
    }

    Ok(())
}